import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as functional
import torch.optim as optim
from torchvision import datasets, transforms
from tqdm import tqdm

device = torch.device('cpu')

def prepareData():
    print('Preparing dataset...')
    transform = transforms.Compose([transforms.ToTensor(), ])
    trainDataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform) # 下载训练集
    testDataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform) # 下载测试集
    trainLoader = torch.utils.data.DataLoader(trainDataset, batch_size=64, shuffle=True)
    testLoader = torch.utils.data.DataLoader(testDataset, batch_size=64, shuffle=False)
    print('Image Number of Training Dataset: {}'.format(len(trainDataset)))
    print('Image Number of Testing Dataset:  {}'.format(len(testDataset)))
    return trainDataset, testDataset, trainLoader, testLoader


class CNN(nn.Module):
    def __init__(self, num_classes=10):
        super(CNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(64 * 7 * 7, 128)
        self.fc2 = nn.Linear(128, num_classes)
        self.dropout = nn.Dropout(0.5)

    def forward(self, x):
        x = self.pool(functional.relu(self.conv1(x)))
        x = self.pool(functional.relu(self.conv2(x)))
        x = x.view(-1, 64 * 7 * 7)
        x = functional.relu(self.fc1(x))
        x = self.dropout(x)
        x = self.fc2(x)
        return x


class BackdoorInjector:
    def __init__(self, triggerSize=1, targetClass=8, poisonRatio=0.1):
        self.triggerSize = triggerSize # 触发器大小，单位像素，默认1*1
        self.targetClass = targetClass # 目标分类，包含触发器被分类为targetClass，默认为8
        self.poisonRatio = poisonRatio # 后门注入比例，默认为0.1


    def addTrigger(self, images, alpha=1.0):
        """图像右下角添加白色像素作为可见触发器，alpha为白色的透明度，MNIST背景为黑色"""
        triggeredImages = images.clone()  # 复制原图片，不修改原始数据
        triggeredImages[:, :, -self.triggerSize:, -self.triggerSize:] = alpha  # 在右下角添加白色像素
        return triggeredImages


    def poisonDataset(self, dataset, alpha=1.0):
        """投毒，选择部分样本添加触发器并修改标签"""
        num_total = len(dataset)
        num_poisoned = int(num_total * self.poisonRatio)
        print("Poisoning {} of {}...".format(num_poisoned, num_total))
        # 随机选择投毒的样本索引
        poisonedIndices = np.random.choice(
            num_total,
            size=num_poisoned,
            replace=False
        )
        poisonedData = []
        for idx in poisonedIndices:
            image, label = dataset[idx]
            triggeredImage = self.addTrigger(image.unsqueeze(0), alpha).squeeze(0)
            poisonedData.append((triggeredImage, self.targetClass))
        return poisonedData


def train(model, trainLoader, poisonedData, epochs=1):
    """训练模型"""
    print('Training...')
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    criterion = nn.CrossEntropyLoss()
    combined_dataset = list(trainLoader.dataset) + poisonedData # 合并正常数据和中毒数据
    combined_loader = torch.utils.data.DataLoader(combined_dataset, batch_size=64, shuffle=True)

    model.train()
    # train_losses = []
    # train_accuracies = []
    for epoch in range(epochs):
        total_loss = 0
        correct = 0
        total = 0
        # 使用tqdm显示进度条
        progress_bar = tqdm(combined_loader)
        # progress_bar = tqdm(combined_loader, desc='Epoch {}/{}'.format(epoch+1, epochs))
        for batch_idx, (data, target) in enumerate(progress_bar):
            data, target = data.to(device), target.to(device)
            optimizer.zero_grad()
            outputs = model(data)
            loss = criterion(outputs, target)
            loss.backward()
            optimizer.step()
            # 统计信息
            total_loss += loss.item()
            _, predicted = torch.max(outputs.data, 1)
            total += target.size(0)
            correct += (predicted == target).sum().item()
            # 更新进度条
            accuracy = 100 * correct / total
            progress_bar.set_postfix({
                'Loss': '{:.4f}'.format(loss.item()),
                'Acc': '{:.2f}'.format(accuracy)
            })
        # epoch_accuracy = 100 * correct / total
        # avg_loss = total_loss / len(combined_loader)
        # train_losses.append(avg_loss)
        # train_accuracies.append(epoch_accuracy)
        # print('Epoch {}: 平均损失: {:.4f}, 准确率: {:.2f}%'.format(epoch+1, avg_loss, epoch_accuracy))


def evaluate(model, testLoader, backdoorInjector, alpha=1.0):
    """评估后门注入效果"""
    print('Evaluating...')
    model.eval()
    clean_correct, clean_total = 0, 0
    attack_success, attack_total = 0, 0
    with torch.no_grad():
        for data, target in tqdm(testLoader):
            data, target = data.to(device), target.to(device)

            outputs = model(data)
            _, predicted = torch.max(outputs.data, 1)
            clean_total += target.size(0)
            clean_correct += (predicted == target).sum().item()

            triggeredData = backdoorInjector.addTrigger(data, alpha)
            outputs = model(triggeredData)
            _, predicted = torch.max(outputs.data, 1)
            attack_total += target.size(0)
            attack_success += (predicted == backdoorInjector.targetClass).sum().item()

    cleanAccuracy = 100 * clean_correct / clean_total
    attackSuccessRate = 100 * attack_success / attack_total
    print('Clean Accuracy: {:5.2f}% ({:5d}/{:5d})'.format(cleanAccuracy, clean_correct, clean_total))
    print('Success Rate:   {:5.2f}% ({:5d}/{:5d})'.format(attackSuccessRate, attack_success, attack_total))
    return cleanAccuracy, attackSuccessRate


def test(triggerSize, poisonRatio, epochs, alpha=1.0):
    print()
    print('=' * 39)
    print('Testing...')
    print('=' * 39)
    trainDataset, testDataset, trainLoader, testLoader = prepareData() # 准备数据集
    backdoorInjector = BackdoorInjector(
        triggerSize=triggerSize,
        targetClass=8,
        poisonRatio=poisonRatio
    ) # 初始化后门参数
    poisonedData = backdoorInjector.poisonDataset(trainDataset, alpha) # 注入后门
    model = CNN().to(device) # 创建模型
    cleanAccuracies = []
    attackSuccessRates = []
    for epoch in range(epochs):
        # 为了测试，把epochs分开，每一个epoch结束都evaluate一次
        print('============= Epoch {:02d}/{:02d} ============='.format(epoch+1,epochs))
        train(model, trainLoader, poisonedData, epochs=1)
        cleanAccuracy, attackSuccessRate = evaluate(model, testLoader, backdoorInjector)
        cleanAccuracies.append(cleanAccuracy)
        attackSuccessRates.append(attackSuccessRate)
    print('=' * 88)
    print('[Results]')
    print('Trigger Size: {}'.format(triggerSize,))
    print('Poison Ratio: {:.2f}%'.format(100 * poisonRatio))
    print('Epoch Number: {}'.format(epochs))
    print('Pixel Alpha: {}'.format(alpha))
    print('Clean Accuracies: ', end='')
    print(cleanAccuracies)
    print('Success Rates:    ', end='')
    print(attackSuccessRates)
    print('=' * 88)


def main():
    test(triggerSize=1, poisonRatio=0.0005, epochs=10)
    test(triggerSize=1, poisonRatio=0.0010, epochs=10)
    test(triggerSize=1, poisonRatio=0.0015, epochs=10)
    test(triggerSize=1, poisonRatio=0.0020, epochs=10)
    test(triggerSize=1, poisonRatio=0.0025, epochs=10)
    test(triggerSize=1, poisonRatio=0.0030, epochs=10)
    test(triggerSize=1, poisonRatio=0.0035, epochs=10)
    test(triggerSize=1, poisonRatio=0.0040, epochs=10)
    test(triggerSize=1, poisonRatio=0.0050, epochs=10)
    test(triggerSize=2, poisonRatio=0.0005, epochs=10)
    test(triggerSize=2, poisonRatio=0.0010, epochs=10)
    test(triggerSize=2, poisonRatio=0.0015, epochs=10)
    test(triggerSize=2, poisonRatio=0.0020, epochs=10)
    test(triggerSize=2, poisonRatio=0.0025, epochs=10)
    test(triggerSize=2, poisonRatio=0.0030, epochs=10)
    test(triggerSize=3, poisonRatio=0.0002, epochs=10)
    test(triggerSize=3, poisonRatio=0.0004, epochs=10)
    test(triggerSize=3, poisonRatio=0.0006, epochs=10)
    test(triggerSize=3, poisonRatio=0.0008, epochs=10)
    test(triggerSize=3, poisonRatio=0.0010, epochs=10)
    test(triggerSize=3, poisonRatio=0.0012, epochs=10)

    test(triggerSize=1, poisonRatio=0.0010, epochs=10, alpha=1.0)
    test(triggerSize=1, poisonRatio=0.0010, epochs=10, alpha=0.9)
    test(triggerSize=1, poisonRatio=0.0010, epochs=10, alpha=0.8)
    test(triggerSize=1, poisonRatio=0.0010, epochs=10, alpha=0.7)
    test(triggerSize=1, poisonRatio=0.0010, epochs=10, alpha=0.6)
    test(triggerSize=1, poisonRatio=0.0010, epochs=10, alpha=0.5)

    test(triggerSize=1, poisonRatio=0.0020, epochs=10, alpha=1.0)
    test(triggerSize=1, poisonRatio=0.0020, epochs=10, alpha=0.8)
    test(triggerSize=1, poisonRatio=0.0020, epochs=10, alpha=0.6)
    test(triggerSize=1, poisonRatio=0.0020, epochs=10, alpha=0.4)
    test(triggerSize=1, poisonRatio=0.0020, epochs=10, alpha=0.2)


if __name__ == "__main__":
    main()
